/*
 * pi3's Linux kernel Runtime Guard
 *
 * Component:
 *  - x86/AMD64 specific code for Exploit Detection
 *
 * Notes:
 *  - None
 *
 * Timeline:
 *  - Created: 18.IV.2019
 *
 * Author:
 *  - Adam 'pi3' Zabrocki (http://pi3.com.pl)
 *
 */

#ifdef CONFIG_X86

#ifndef P_LKRG_EXPLOIT_DETECTION_X86_ARCH_H
#define P_LKRG_EXPLOIT_DETECTION_X86_ARCH_H

/*
 * pCFI global CPU flags
 */
#define P_PCFI_X86_WP     0x1
#define P_PCFI_X86_SMEP   0x2
#define P_PCFI_X86_SMAP   0x3

#define P_IS_WP_FLAG_ENABLED(x)   (x & P_PCFI_X86_WP)
#define P_IS_SMEP_FLAG_ENABLED(x) (x & P_PCFI_X86_SMEP)
#define P_IS_SMAP_FLAG_ENABLED(x) (x & P_PCFI_X86_SMAP)

/*
 * Known minor bug: these updates are non-atomic, so concurrent changing of
 * these flags might result in all but one of the changes getting lost.  We
 * may switch to using individual Boolean variables instead of the bitmask.
 */
#define P_ENABLE_WP_FLAG(x)       (x |= P_PCFI_X86_WP)
#define P_ENABLE_SMEP_FLAG(x)     (x |= P_PCFI_X86_SMEP)
#define P_ENABLE_SMAP_FLAG(x)     (x |= P_PCFI_X86_SMAP)

#define P_DISABLE_WP_FLAG(x)      (x &= ~P_PCFI_X86_WP)
#define P_DISABLE_SMEP_FLAG(x)    (x &= ~P_PCFI_X86_SMEP)
#define P_DISABLE_SMAP_FLAG(x)    (x &= ~P_PCFI_X86_SMAP)

/*
 * x86 specific functions
 */
static inline void p_write_cr4(unsigned long p_arg) {

#if LINUX_VERSION_CODE >= KERNEL_VERSION(5,8,0)
   P_SYM(p_native_write_cr4)(p_arg);
#elif LINUX_VERSION_CODE < KERNEL_VERSION(4,0,0)
   write_cr4(p_arg);
#else
   __write_cr4(p_arg);
#endif

}

static inline unsigned long p_read_cr4(void) {

#if LINUX_VERSION_CODE < KERNEL_VERSION(4,0,0)
   return read_cr4();
#else
   return __read_cr4();
#endif

}

static inline unsigned int p_ed_pcfi_x86_validate_wp(unsigned char p_kill) {

   unsigned int p_bad = 0;
   register unsigned long p_cr0;

   if (!P_CTRL(p_kint_validate))
      return p_bad;

   p_cr0 = read_cr0();

   if (!(p_cr0 & X86_CR0_WP)) {

      switch (P_CTRL(p_kint_enforce)) {

         /* Panic */
         case 2:
           // OK, we need to crash the kernel now
           panic(P_LKRG_SIGNATURE "Write Protection bit was disabled and 'panic()' is enforced! Killing the kernel...\n");
           break;

         /* Log and restore */
         case 1:
            p_print_log(P_LKRG_CRIT,
                   "Write Protection bit was disabled! Enforcing WP now!\n");
            p_cr0 |= X86_CR0_WP;
            write_cr0(p_cr0);
            wbinvd();
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_WP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LKRG_CRIT,
                   "Write Protection bit was disabled! Accepting new WP state.\n");
            P_DISABLE_WP_FLAG(p_pcfi_CPU_flags);
            break;
      }

      /*
       * It is possible to kill an almost arbitrary process here. Such situation is possible
       * if process which disables SMEP / WP is being preempted before leverages such state.
       * Most of the exploits which I tested didn't have such FP, but such possibility of course
       * exists. However, it is worth the risk. Especially that such changes (SMEP / WP violation)
       * are not legal and are always the results of malicious activity.
       */
      if (p_kill) {
         if (P_CTRL(p_pint_enforce)) {
            p_print_log(P_LKRG_CRIT, "We might be killing arbitrary process here. However, it is worth the risk!\n");
            p_ed_kill_task_by_task(current);
         }
      }
   }

   return p_bad;
}

static inline unsigned int p_ed_pcfi_x86_validate_smXp(unsigned char p_kill) {

   unsigned int p_bad = 0;
   register unsigned long p_cr4;


   if (!P_CTRL(p_smep_validate) && !P_CTRL(p_smap_validate))
      return p_bad;

   p_cr4 = p_read_cr4();

   if (!(p_cr4 & X86_CR4_SMEP) && P_CTRL(p_smep_validate)) {

      switch (P_CTRL(p_smep_enforce)) {

         /* Panic */
         case 2:
           // OK, we need to crash the kernel now
           panic(P_LKRG_SIGNATURE "SMEP was disabled and 'panic()' is enforced! Killing the kernel...\n");
           break;

         /* Log and restore */
         case 1:
            p_print_log(P_LKRG_CRIT,
                   "SMEP was disabled! Enforcing SMEP now!\n");
//            cr4_set_bits(X86_CR4_SMEP);
            p_cr4 |= X86_CR4_SMEP;
            p_write_cr4(p_cr4);
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_SMEP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LKRG_CRIT,
                   "SMEP was disabled! Accepting new SMEP state.\n");
            P_DISABLE_SMEP_FLAG(p_pcfi_CPU_flags);
            break;
      }

      /*
       * It is possible to kill an almost arbitrary process here. Such situation is possible
       * if process which disables WP / SMEP / SMAP is being preempted before leverages such state.
       * Most of the exploits which I tested didn't have such FP, but such possibility of course
       * exists. However, it is worth the risk. Especially that such changes (WP / SMEP / SMAP violation)
       * are not legal and are always the results of malicious activity.
       */
      if (p_kill) {
         if (P_CTRL(p_pint_enforce)) {
            p_print_log(P_LKRG_CRIT, "We might be killing arbitrary process here. However, it is worth the risk!\n");
            p_ed_kill_task_by_task(current);
         }
      }
   }

   if (!(p_cr4 & X86_CR4_SMAP) && P_CTRL(p_smap_validate)) {

      switch (P_CTRL(p_smap_enforce)) {

         /* Panic */
         case 2:
           // OK, we need to crash the kernel now
           panic(P_LKRG_SIGNATURE "SMAP was disabled and 'panic()' is enforced! Killing the kernel...\n");
           break;

         /* Log and restore */
         case 1:
            p_print_log(P_LKRG_CRIT,
                   "SMAP was disabled! Enforcing SMAP now!\n");
//            cr4_set_bits(X86_CR4_SMAP);
            p_cr4 |= X86_CR4_SMAP;
            p_write_cr4(p_cr4);
            p_bad++;
            break;

         /* Log and accept */
         case 0:
            if (!P_IS_SMAP_FLAG_ENABLED(p_pcfi_CPU_flags))
               return p_bad;
            p_print_log(P_LKRG_CRIT,
                   "SMAP was disabled! Accepting new SMAP state.\n");
            P_DISABLE_SMAP_FLAG(p_pcfi_CPU_flags);
            break;
      }

      /*
       * It is possible to kill an almost arbitrary process here. Such situation is possible
       * if process which disables WP / SMEP / SMAP is being preempted before leverages such state.
       * Most of the exploits which I tested didn't have such FP, but such possibility of course
       * exists. However, it is worth the risk. Especially that such changes (WP / SMEP / SMAP violation)
       * are not legal and are always the results of malicious activity.
       */
      if (p_kill) {
         if (P_CTRL(p_pint_enforce)) {
            p_print_log(P_LKRG_CRIT, "We might be killing arbitrary process here. However, it is worth the risk!\n");
            p_ed_kill_task_by_task(current);
         }
      }
   }

   return p_bad;
}

#endif

#endif
